//
//  RSAPublicKeyEncryptor.m
//  Tunnel
//
//  Created by wbb on 2025/6/9.
//

#import "RSAPublicKeyEncryptor.h"
#import <Security/Security.h>
#import <CommonCrypto/CommonCrypto.h>

@interface RSAPublicKeyEncryptor ()

@property (nonatomic, assign) SecKeyRef publicKey;

@end

@implementation RSAPublicKeyEncryptor

- (instancetype)initWithPublicKeyPEM:(NSString *)publicKeyPEM {
    self = [super init];
    if (self) {
        // 移除PEM头和尾
        NSString *keyString = [publicKeyPEM stringByReplacingOccurrencesOfString:@"-----BEGIN PUBLIC KEY-----" withString:@""];
        keyString = [keyString stringByReplacingOccurrencesOfString:@"-----END PUBLIC KEY-----" withString:@""];
        keyString = [keyString stringByTrimmingCharactersInSet:[NSCharacterSet whitespaceAndNewlineCharacterSet]];
        
        // 转换为NSData
        NSData *publicKeyData = [[NSData alloc] initWithBase64EncodedString:keyString options:0];
        if (publicKeyData) {
            [self loadPublicKeyWithData:publicKeyData];
        }
    }
    return self;
}

- (instancetype)initWithPublicKeyBase64:(NSString *)publicKeyBase64 {
    self = [super init];
    if (self) {
        NSData *publicKeyData = [[NSData alloc] initWithBase64EncodedString:publicKeyBase64 options:0];
        if (publicKeyData) {
            [self loadPublicKeyWithData:publicKeyData];
        }
    }
    return self;
}

- (void)dealloc {
    if (_publicKey) {
        CFRelease(_publicKey);
        _publicKey = NULL;
    }
}

- (void)loadPublicKeyWithData:(NSData *)publicKeyData {
    // 创建公钥
    NSDictionary *attributes = @{
        (id)kSecAttrKeyType: (id)kSecAttrKeyTypeRSA,
        (id)kSecAttrKeyClass: (id)kSecAttrKeyClassPublic,
        (id)kSecReturnPersistentRef: @YES
    };
    
    CFErrorRef error = NULL;
    self.publicKey = SecKeyCreateWithData((CFDataRef)publicKeyData, (CFDictionaryRef)attributes, &error);
    
    if (error) {
        NSLog(@"加载RSA公钥失败: %@", CFBridgingRelease(error));
    }
}

- (NSString *)encryptString:(NSString *)plainText {
    if (!self.publicKey || !plainText) {
        return nil;
    }
    
    NSData *plainData = [plainText dataUsingEncoding:NSUTF8StringEncoding];
    NSData *encryptedData = [self encryptData:plainData];
    
    if (encryptedData) {
        return [encryptedData base64EncodedStringWithOptions:0];
    }
    
    return nil;
}

- (NSData *)encryptData:(NSData *)plainData {
    if (!self.publicKey || !plainData) {
        return nil;
    }
    
    size_t keySize = SecKeyGetBlockSize(self.publicKey);
    size_t bufferSize = keySize - 11; // PKCS#1 padding
    
    NSMutableData *encryptedData = [NSMutableData data];
    const uint8_t *srcBytes = [plainData bytes];
    size_t srcLength = [plainData length];
    
    for (size_t i = 0; i < srcLength; i += bufferSize) {
        size_t length = MIN(bufferSize, srcLength - i);
        uint8_t *buffer = malloc(keySize);
        size_t encryptedLength = keySize;
        
        OSStatus status = SecKeyEncrypt(self.publicKey,
                                        kSecPaddingPKCS1,
                                        &srcBytes[i],
                                        length,
                                        buffer,
                                        &encryptedLength);
        
        if (status == noErr) {
            [encryptedData appendBytes:buffer length:encryptedLength];
        } else {
            NSLog(@"RSA加密失败: %d", (int)status);
            free(buffer);
            return nil;
        }
        
        free(buffer);
    }
    
    return encryptedData;
}


@end
